import {
  SubscribeMessage,
  WebSocketGateway,
  WebSocketServer,
  MessageBody,
  ConnectedSocket,
  OnGatewayConnection,
  OnGatewayDisconnect
} from '@nestjs/websockets';
import { Server, Socket } from 'socket.io';
import { Injectable, Logger } from '@nestjs/common';
import { JwtService } from '@nestjs/jwt';
import { MessageService } from '../../message/message.service';

interface AuthenticatedSocket extends Socket {
  userId?: string;
  sessionId?: string;
  personaId?: string;
  skill?: string;
}

@Injectable()
@WebSocketGateway({
  cors: { origin: '*', methods: ['GET','POST'] },
  namespace: '/chat'
})
export class WsGateway implements OnGatewayConnection, OnGatewayDisconnect {
  @WebSocketServer()
  server: Server;

  private readonly logger = new Logger(WsGateway.name);
  private connectedClients = new Map<string, AuthenticatedSocket>();

  constructor(
    private jwtService: JwtService,
    private messageService: MessageService,
  ) {}

  // Handle client connection
  async handleConnection(client: AuthenticatedSocket) {
    this.logger.log(`Client connected: ${client.id}`);
  }

  // Handle client disconnection
  handleDisconnect(client: AuthenticatedSocket) {
    this.logger.log(`Client disconnected: ${client.id}`);
    if (client.sessionId) {
      this.connectedClients.delete(client.sessionId);
    }
  }

  // Handle user authentication
  @SubscribeMessage('authenticate')
  async handleAuthenticate(
    @MessageBody() data: { accessToken: string; sessionId: string },
    @ConnectedSocket() client: AuthenticatedSocket,
  ) {
    try {
      // Verify accessToken
      const userPayload = this.jwtService.verify(data.accessToken);
      client.userId = userPayload.sub;
      client.sessionId = data.sessionId;
      
      // Join session room
      await client.join(`session:${data.sessionId}`);
      this.connectedClients.set(data.sessionId, client);

      return { success: true, sessionId: data.sessionId };
    } catch (error) {
      this.logger.error(`Authentication failed for client ${client.id}:`, error);
      return { success: false, error: 'Invalid access token' };
    }
  }

  // Handle send message
  @SubscribeMessage('send_message')
  async handleSendMessage(
    @MessageBody() data: { content: string; metadata?: any },
    @ConnectedSocket() client: AuthenticatedSocket,
  ) {
    if (!client.sessionId || !client.userId) {
      return { success: false, error: 'Not authenticated' };
    }

    try {
      // Save message to database
      const response = await this.messageService.create(client.userId, {
        content: data.content,
        sessionId: client.sessionId,
        metadata: data.metadata,
      });

      // Broadcast user message to session room
      this.server.to(`session:${client.sessionId}`).emit('message_received', {
        id: response.userMessage.id,
        role: 'USER',
        content: response.userMessage.content,
        metadata: response.userMessage.metadata,
        createdAt: response.userMessage.createdAt,
        userId: client.userId,
      });

      // Broadcast AI response to session room
      this.server.to(`session:${client.sessionId}`).emit('message_received', {
        id: response.aiResponse.id,
        role: 'ASSISTANT',
        content: response.aiResponse.content,
        metadata: response.aiResponse.metadata,
        createdAt: response.aiResponse.createdAt,
      });

      return { success: true };
    } catch (error) {
      this.logger.error(`Error sending message:`, error);
      return { success: false, error: 'Failed to send message' };
    }
  }

  // Handle voice chunk
  @SubscribeMessage('voice_chunk')
  handleVoiceChunk(
    @MessageBody() data: { sessionId: string; bytes: ArrayBuffer; personaId?: string; skill?: string; meta?: any },
    @ConnectedSocket() client: AuthenticatedSocket,
  ) {
    if (!client.sessionId) {
      return { success: false, error: 'Not authenticated' };
    }

    // Broadcast voice data to session room
    this.server.to(`session:${client.sessionId}`).emit('tts_chunk', { 
      bytes: data.bytes,
      sessionId: client.sessionId 
    });
    
    return { success: true };
  }

  // Handle end of utterance
  @SubscribeMessage('end_of_utterance')
  handleEoU(@MessageBody() body: { sessionId: string }, @ConnectedSocket() client: AuthenticatedSocket) {
    if (!client.sessionId) {
      return { success: false, error: 'Not authenticated' };
    }

    this.server.to(`session:${client.sessionId}`).emit('utterance_ended', {
      sessionId: client.sessionId,
      timestamp: new Date().toISOString()
    });

    return { success: true };
  }

  // Get online user count
  @SubscribeMessage('get_online_count')
  async handleGetOnlineCount(@ConnectedSocket() client: AuthenticatedSocket) {
    if (!client.sessionId) {
      return { success: false, error: 'Not authenticated' };
    }

    const room = this.server.sockets.adapter.rooms.get(`session:${client.sessionId}`);
    const count = room ? room.size : 0;
    
    return { success: true, count };
  }
}